import json
import os
import random  # 新增：用于打乱数据

def process_json_file(input_file, output_dir, shuffle=True):  # 新增shuffle参数
    """
    处理JSON文件，提取指定字段并按dataset分类保存（可选打乱）
    :param shuffle: 是否打乱每个数据集内的数据顺序（默认True）
    """
    try:
        with open(input_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        datasets = ['train', 'val', 'test']
        for dataset in datasets:
            os.makedirs(os.path.join(output_dir, dataset), exist_ok=True)
        
        # 初始化：按数据集分组存储（先集中收集，再统一处理）
        grouped_data = {dataset: [] for dataset in datasets}
        prompt = "描述图像中的空间布局和设计细节。"
        
        # 1. 先按dataset分组收集所有数据
        for item in data:
            dataset = item.get("dataset", "train")
            if dataset not in datasets:
                dataset = "train"
            tag = item.get("tag", "").replace("|", ",")
            processed_item = {
                "image": item.get("image", ""),
                "prompt": prompt,
                "tag": tag
            }
            grouped_data[dataset].append(processed_item)
        
        # 2. 可选：打乱每个数据集内的数据顺序
        if shuffle:
            for dataset in datasets:
                random.shuffle(grouped_data[dataset])  # 打乱当前数据集内的顺序
        
        # 3. 写入文件
        stats = {}
        for dataset in datasets:
            output_file = os.path.join(output_dir, dataset, f"data_tag_{dataset}.json")
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(grouped_data[dataset], f, ensure_ascii=False, indent=2)
            stats[dataset] = len(grouped_data[dataset])
        
        print(f"\n处理完成！数据分布统计：")
        for dataset, count in stats.items():
            print(f"{dataset}集: {count}条 → {output_file}")
    
    except Exception as e:
        print(f"处理过程中发生错误：{str(e)}")

# 使用示例
if __name__ == "__main__":
    input_json = "final_dataset_GZ.json"
    output_dir = "processed_NewGZ_tag"
    process_json_file(input_json, output_dir, shuffle=True)  # 启用打乱